import os
import json
import numpy as np
import faiss
import requests
from typing import List, Dict, Tuple

# 已有的文本转向量函数
def get_embedding(text: str) -> np.ndarray:
    """
    使用在线API获取文本的向量表示
    """
    url = "https://api.siliconflow.cn/v1/embeddings"

    payload = {
        "model": "BAAI/bge-m3",
        "input": text,
        "encoding_format": "float"
    }
    headers = {
        "Authorization": "Bearer sk-wdotugwguprsmeimsowcbehipjlkyaabkgbxnudkkhvjhumm",
        "Content-Type": "application/json"
    }

    response = requests.post(url, json=payload, headers=headers)
    response.raise_for_status()  # 检查请求是否成功
    
    embedding = np.array(response.json()["data"][0]["embedding"])
    return embedding

class RAGSystem:
    def __init__(self, index_file: str = "faiss_index.faiss", metadata_file: str = "metadata.json"):
        """
        初始化RAG系统
        
        参数:
            index_file: FAISS索引文件路径
            metadata_file: 元数据文件路径
        """
        self.index_file = index_file
        self.metadata_file = metadata_file
        self.index = None
        self.metadata = []
        
        # 加载索引和元数据
        self._load_index()
    
    def _load_index(self):
        """
        加载FAISS索引和元数据
        """
        try:
            if os.path.exists(self.index_file):
                print("加载现有的FAISS索引...")
                self.index = faiss.read_index(self.index_file)
            
            if os.path.exists(self.metadata_file):
                print("加载现有的元数据...")
                with open(self.metadata_file, 'r', encoding='utf-8') as f:
                    self.metadata = json.load(f)
            
            # 检查索引和元数据是否匹配
            if self.index is not None and self.metadata:
                if len(self.metadata) != self.index.ntotal:
                    print("警告: 索引和元数据数量不匹配")
        
        except Exception as e:
            print(f"加载索引失败: {e}")
            self.index = None
            self.metadata = []
    
    def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
        """
        搜索最相似的文档
        
        参数:
            query: 用户查询
            top_k: 返回最相似的top_k个结果
            
        返回:
            列表: [(文档内容, 相似度得分), ...]
        """
        if self.index is None:
            raise RuntimeError("索引未加载，请确保已正确构建数据库")
        
        # 将查询转化为向量
        query_embedding = get_embedding(query).reshape(1, -1)
        
        # 搜索相似向量
        distances, indices = self.index.search(query_embedding, top_k)
        
        # 获取结果
        results = []
        for i, idx in enumerate(indices[0]):
            if idx == -1 or idx >= len(self.metadata):  # 如果没有找到足够的匹配项或索引越界
                continue
            doc_info = self.metadata[idx]
            results.append((doc_info["text"], distances[0][i]))
        
        return results
    
    def answer_question(self, question: str, top_k: int = 3) -> str:
        """
        回答用户问题
        
        参数:
            question: 用户问题
            top_k: 考虑最相似的top_k个文档
            
        返回:
            str: 回答内容
        """
        if self.index is None:
            return "错误：知识库尚未初始化，请先构建数据库"
        
        # 搜索最相似的文档
        results = self.search(question, top_k)
        
        if not results:
            return "抱歉，没有找到相关答案。"
        
        # 这里简单地将所有相关文档内容拼接作为回答
        # 实际应用中可以使用更复杂的答案提取方法
        answer = "根据知识库内容:\n\n"
        # 按照 score 从大到小排序
        sorted_results = sorted(results, key=lambda x: x[1], reverse=True)

        for i, (text, score) in enumerate(sorted_results, 1):
            # 提取简短摘要而不是整个文档
            # 这里简单截取前200个字符，实际可以更智能
            answer += f"结果 {i} (相似度: {score:.4f}):\n{text}\n\n"
        
        return answer.strip()

if __name__ == "__main__":
    # 初始化RAG系统
    rag = RAGSystem()
    
    # 测试问答
    print("RAG问答系统已启动 (输入'quit'退出)")
    while True:
        question = input("\n请输入你的问题: ")
        if question.lower() == 'quit':
            break
        answer = rag.answer_question(question)
        print("\n回答:")
        print(answer)